# -*- coding: utf-8 -*-
"""
Created on Sat Oct 20 19:45:03 2018

@author: jia.liu
"""
import pandas as pd
from my_data_describe import my_data_describe
from sklearn.linear_model import LinearRegression
from sklearn.linear_model import LogisticRegression

def my_fillna_model(df, narate=0.4, reg_class_split_num = 20):
    '''
    根据其它列，预测缺失列（与y无关）
    narate：缺失率在此之下，采用模型处理
    reg_class_split_num：判断用回归or分类的种类数量
    '''
    train_test = df.copy()
    df_describe = my_data_describe(train_test)
    ### 缺失较少的 回归填充
    na_little = df_describe[(df_describe.na_rate<=narate) & (df_describe.na_rate>0)]
    na_reg = na_little[(na_little.type=='float64')&(na_little.unique_num>reg_class_split_num)]
    na_class = na_little[(na_little.type=='object')|(na_little.unique_num<=reg_class_split_num)]
    reg_li = na_reg.index
    class_li = na_class.index
    all_col = len(class_li)
    done_col = 0
    # 类别型填充
    for class_col in class_li:
        temp_others = train_test.drop(class_col,axis=1)
        temp_others = pd.get_dummies(temp_others)
        temp_others = temp_others.fillna(temp_others.mean())
        temp_test_index = train_test[class_col][train_test[class_col].isnull()].index
        temp_test = temp_others.iloc[temp_test_index]
        temp_train_X = temp_others.drop(temp_test_index)
        temp_train_y = train_test[class_col].drop(temp_test_index)
        train_test.drop(class_col,axis=1)
        lr = LogisticRegression()
        lr.fit(temp_train_X,temp_train_y)
        train_test[class_col][temp_test_index] = lr.predict(temp_test)
        train_test[class_col][temp_test_index]
        done_col += 1
        print('%d/%d\t%s\tfilled'%(done_col,all_col,class_col))
    all_col = len(reg_li)
    done_col = 0
    for reg_col in reg_li:
        temp_others = train_test.drop(reg_col,axis=1)
        temp_others = pd.get_dummies(temp_others)
        temp_others = temp_others.fillna(temp_others.mean())
        temp_test_index = train_test[reg_col][train_test[reg_col].isnull()].index
        temp_test = temp_others.iloc[temp_test_index]
        temp_train_X = temp_others.drop(temp_test_index)
        temp_train_y = train_test[reg_col].drop(temp_test_index)
        train_test.drop(reg_col,axis=1)
        linReg = LinearRegression()
        linReg.fit(temp_train_X,temp_train_y)
        max_num = temp_train_y.max()
        min_num = temp_train_y.min()
        predict_li = list(linReg.predict(temp_test))
        def max_min_cut(x):
            if x >= max_num:
                return max_num
            elif x <= min_num:
                return min_num
            else:
                return x
        predict_li_new = list(map(max_min_cut, predict_li))
        train_test[reg_col][temp_test_index] = predict_li_new
        done_col += 1
        print('%d/%d\t%s\tfilled'%(done_col,all_col,reg_col))    
    
    return train_test
